#%%
import numpy as np
import math
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from scipy.signal import convolve2d
from Chaikin import *


def sigmoid(x):
    # print(x,type(x),-x/5+12)
    return 0.99-0.98/ (1 + math.exp(-x/20+6))
def speed_function(gray_level):
    # 这里使用一个简单的线性速度函数，速度与灰度级成正比
    return (sigmoid(gray_level))
def speed_function2(gray_level):
    if gray_level>120:
        return 0.1
    else:
        return 1


#%%

def wavefront_propagation(image, start_point,end_point,end_point2=None):
    rows, cols = image.shape
    arrival_times = np.inf * np.ones((rows, cols))
    arrival_times[start_point[0],start_point[1]] = 0
    wavefront = [start_point]
    directions = [(-1, 0,1), (1, 0,1), (0, -1,1), (0, 1,1), (-1, -1,1.414), (-1, 1,1.414), (1, -1,1.414), (1, 1,1.414)]  # 八连通方向
    stop1=False 
    stop2=False
    if end_point2==None:
        stop2=True
    while wavefront and (not stop1 or not stop2):
        x, y = wavefront.pop(0)
        for dx, dy,distance in directions:
            nx, ny = x + dx, y + dy
            if 0 <= nx < rows and 0 <= ny < cols:
                # 计算距离，这里假设相邻点距离为 1
                # distance = 1
                
                speed = speed_function(int(image[nx, ny]))
                new_arrival_time = arrival_times[x, y] + distance / speed
                if new_arrival_time < arrival_times[nx, ny]:
                    arrival_times[nx, ny] = new_arrival_time
                    wavefront.append((nx, ny))
                    # print((nx,ny),end_point)
                    if nx==end_point[0] and ny==end_point[1]:
                        print('stop')
                        stop1=True
                    if end_point2!=None and nx==end_point2[0] and ny==end_point2[1]:
                        print('stop')
                        stop2=True

    return arrival_times
#%%


def backtracking(arrival_times, endpoint):
    """
    回溯法找到从终点到起点的路径。
    
    参数:
    arrival_times: 一个二维数组，记录每个位置的到达时间。
    endpoint: 元组，表示终点的坐标。
    
    返回:
    一条从终点到起点的路径，表示为一个坐标列表。
    """
    # 获取arrival_times的行数和列数
    rows, cols = arrival_times.shape
    # 初始化路径列表，起点为终点
    path = [endpoint]
    # 定义八个可能的移动方向，分别表示上、下、左、右、左上、左下、右上、右下
    directions = [(-1, 0), (1, 0), (0, -1), (0, 1), (-1, -1), (-1, 1), (1, -1), (1, 1)]
    # 初始化当前位置为终点
    x, y = endpoint
    # 循环直到无法找到更早的到达时间或回到起点
    while True:
        # 初始化最小时间为无穷大，用于后续比较
        min_time = np.inf
        # 初始化最小时间对应的方向为None
        min_direction = None
        # 遍历所有可能的方向
        for dx, dy in directions:
            # 计算新的位置
            nx, ny = x + dx, y + dy
            # 检查新位置是否在合法范围内
            if 0 <= nx < rows and 0 <= ny < cols:
                # 如果新位置的到达时间小于当前最小时间，则更新最小时间和方向
                if arrival_times[nx, ny] < min_time:
                    min_time = arrival_times[nx, ny]
                    min_direction = (dx, dy)
        # 如果没有找到更早的到达时间，跳出循环
        if min_direction is None:
            break
        # 根据最小时间对应的方向更新当前位置
        x += min_direction[0]
        y += min_direction[1]
        # 将当前位置添加到路径中
        path.append((x, y))
        # 如果到达起始点，跳出循环
        if arrival_times[x, y] == 0:
            break
    # 反转路径，使其从起点到终点，并返回
    return path[::-1]

def edgeStrength(image,path,sampleDistance=40,width=40):
    strengthValue = np.zeros((len(path),width*2+1,3))
    resampledImg = np.zeros((len(path),width*2+1,3))
    # perpendicular line
    perpendicularLines=[]
    step=1
    indexset = list(range(0,len(path),step))
    if indexset[len(indexset)-1]!=len(path)-1:
        indexset.append(len(path)-1)
    print(indexset)
    for index in indexset:
        point2 = path[index-sampleDistance] if index+sampleDistance>=len(path) else path[index+sampleDistance]
        vec2= np.array(path[index])-np.array(point2) if index+sampleDistance>=len(path) else (-np.array(path[index])+np.array(point2) )
        if np.linalg.norm(vec2)!=0:
            vec2 = vec2/np.linalg.norm(vec2)
        # print(index,[np.array(path[index]),np.array([-vec2[1],vec2[0]])])
        for j in range(-width,width+1):
            coord=np.array(path[index])+np.array([-vec2[1],vec2[0]])*j
            resampledImg[index,j+width]=[image[int(coord[0]),int(coord[1])],int(coord[0]),int(coord[1])]
            strengthValue[index,j+width]=[image[int(coord[0]),int(coord[1])],(coord[0]),(coord[1])]
            pass
        perpendicularLines.append([np.array(path[index]),np.array([-vec2[1],vec2[0]])])
    
    # edge filter
    kernel = np.array([[1, -1]])
    convolve2dimage1st = (convolve2d(resampledImg[:,:,0], kernel, mode='same'))
    convolve2dimage1st[:,0]=0
    convolve2dimage2st = (convolve2d(convolve2dimage1st, kernel, mode='same'))
    strengthValue[:,:,0] = convolve2dimage1st+convolve2dimage2st

    return perpendicularLines,strengthValue,resampledImg,convolve2dimage1st,convolve2dimage2st


def dynamic_programming(edge_strengths):
    n = len(edge_strengths)
    m = len(edge_strengths[0])
    cost_matrix = np.zeros((n, m))
    path_matrix = np.zeros((n, m), dtype=int)

    # 初始化第一行
    cost_matrix[0, :] = edge_strengths[0]

    for i in range(1, n):
        for j in range(m):
            min_cost = float('inf')
            min_idx = -1
            for k in range(max(0, j-2), min(m, j+3)):
                cost = cost_matrix[i-1, k] + edge_strengths[i][j]
                if cost < min_cost:
                    min_cost = cost
                    min_idx = k
            cost_matrix[i, j] = min_cost
            path_matrix[i, j] = min_idx

    # 回溯找到最优路径
    min_j = int(np.average(np.argmin(cost_matrix[-50:],axis=1)))
    print(min_j)
    path = [min_j]

    for i in range(n-1, 0, -1):
        min_j = path_matrix[i, min_j]
        path.append(min_j)
    for i in range(1,0,n):
        pass
    path.reverse()

    return path

def mapCountour2img(strengthValue,path):
    countour=[]

    for i in range(len(path)):
        xindex = strengthValue[:,:,1][i,path[i]]
        yindex = strengthValue[:,:,2][i,path[i]]
        countour.append([xindex,yindex])

        # gray_img1[int(xindex),int(yindex)]=255
    return countour
def mca(gray_img,path,smoothStep=15):
    perpendicularLines,strengthValue,resampledImg,convolve2dimage1st,convolve2dimage2st = edgeStrength(gray_img,path,20,40)
    shape = strengthValue.shape
    path1 = dynamic_programming(strengthValue[:,0:int(shape[1]/2),0])
    path2 = dynamic_programming(-strengthValue[:,int(shape[1]/2):,0])
    countour1=mapCountour2img(strengthValue[:,0:int(shape[1]/2),:],path1)
    countour2=mapCountour2img(strengthValue[:,int(shape[1]/2):,:],path2)
    points=[]
    for p in countour1[::smoothStep]:
        points.append(PointD(p[0],p[1]))
    points.append(PointD(countour1[-1][0],countour1[-1][1]))
    smoothed_points = chaikin_smoothing(points, tension=0.5, nr_of_iterations=3)
    points=[]
    for p in countour2[::smoothStep]:
        points.append(PointD(p[0],p[1]))
    points.append(PointD(countour2[-1][0],countour2[-1][1]))
    
    smoothed_points2 = chaikin_smoothing(points, tension=0.5, nr_of_iterations=3)
    return smoothed_points,smoothed_points2,perpendicularLines
# %% test speed function


# %%
if __name__ == "__main__":
    pass
    #%%
    sigmoid_inputs = np.arange(0,170,0.01)

    sigmoid_outputs = []
    for p in sigmoid_inputs:
        sigmoid_outputs.append(speed_function(p))
    print("Sigmoid Function Input :: {}".format(sigmoid_inputs))
    print("Sigmoid Function Output :: {}".format(sigmoid_outputs))

    plt.plot(sigmoid_inputs,sigmoid_outputs)
    plt.xlabel("Sigmoid Inputs")
    plt.ylabel("Sigmoid Outputs")
    plt.show()
# %%
